from langchain.tools import BaseTool
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel,Field,ConfigDict
from typing import Any,Type
from dotenv import load_dotenv
from model.my_chat_model import ChatModel

#定义输入参数的数据模型类型，... 表示必填
class ChatInput(BaseModel):
    question:str = Field(...,description="问题")

#定义工具类
class ChatTool(BaseTool):
    # 定义模型，是否允许输入参数，输入参数的数据模型类
    model_config = ConfigDict(arbitrary_types_allowed=True)
    def __init__(self,**kwargs:Any):
        super().__init__(
            name = "get_chat_tool",
            description="主要用于回答其他问题，输入的参数必须是问题",
            **kwargs
        )
        #加载环境变量
        load_dotenv()
    #定义工具参数
    args_schema : Type[BaseModel] = ChatInput

    #定义工具方法
    def _run(self,question:str):
        #获取模型
        chat=ChatModel()
        llm=chat.get_line_model()
        rs=llm.invoke(question)
        return rs.content

if __name__ == '__main__':
    chat = ChatTool()
    question="西游记是什么"
    res=chat._run(question)
    print(res)